package ws

import (
	// "fmt"
	"net/http"
	"golang.org/x/net/websocket"
	"sync"
	"time"
	"strings"
	"runtime/debug"

	"gitee.com/wint/tge"
)

type WebSocketServer struct {
	Conns		map[string]*websocket.Conn
	CreateAgent	func(tge.Conn) tge.Agent

	server 		*http.Server

	IsStarted	bool
	mu 			sync.RWMutex
	wg 			sync.WaitGroup
}

func NewWebSocketServer() *WebSocketServer {
	s := &WebSocketServer{}
	s.Conns		= make(map[string]*websocket.Conn)
	s.server = &http.Server{}
	s.IsStarted = false
	return s
}

func (s *WebSocketServer) SetCreateAgent (fn func(tge.Conn) tge.Agent) {
	s.CreateAgent = fn
}

func (s *WebSocketServer) wsConnHandler(conn *websocket.Conn) {
	defer func() {
		if err := recover(); err != nil {
			tge.Logger.Println("wsConnHandler", err, string(debug.Stack()))
		}
	}()

	key := conn.RemoteAddr().String()

	s.mu.Lock()
	s.Conns[key] = conn
	s.mu.Unlock()

	defer conn.Close()


	s.wg.Add(1)
	defer s.wg.Done()

	agent := s.CreateAgent(NewWebSocketConn(conn))
	defer agent.Close()
	agent.Run()

	s.mu.Lock()
	delete(s.Conns, key)
	s.mu.Unlock()
}

func (s *WebSocketServer) StartServer(addr string) {
	h := websocket.Handler(s.wsConnHandler)
	svrMux := http.NewServeMux()
	svrMux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Access-Control-Allow-Credentials", "true")
		w.Header().Add("Access-Control-Allow-Origin", "*")
		h.ServeHTTP(w, r)
	})
	if !strings.Contains(addr, ":") {
		addr = ":" + addr
	}
	tge.Logger.Println("Listen Websocket on", addr)
	s.IsStarted = true
	s.server.Addr = addr
	s.server.Handler = svrMux
	s.server.ListenAndServe()
}

func (s *WebSocketServer) Close() {
	if !s.IsStarted {
		tge.Logger.Println("WebSocketServer isn't started")
		return
	}

	s.IsStarted = false

	s.server.Close()

	t := time.Now().Add(100 * time.Millisecond)
	for _, v := range s.Conns {
		v.SetDeadline(t)
	}

	tge.Logger.Println("WebSocketServer waiting for threads end")
	s.wg.Wait()

	tge.Logger.Println("WebSocketServer is closed")
}
